import matplotlib.pyplot as plt
import numpy as np
import rasterio
import os
import geopandas as gpd
from rasterio.mask import mask
from matplotlib import rcParams
import matplotlib.font_manager as fm

# 设置中文字体
rcParams['font.sans-serif'] = ['Noto Sans CJK JP']  # 选择简体中文
rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

pwd = os.getcwd()

# 定义掩膜文件和对应的中文名称
masks = [
    ('shannan_1', '陕南'),
    ('guanzhong_1', '关中'),
    ('shanbei_1', '陕北')
]

# 初始化存储结果的字典
results = {}

# 加载掩膜数据并计算 NPP 均值
for mask_name, mask_chinese in masks:
    shapefile_path = '{}.shp'.format(mask_name)
    mask_gdf = gpd.read_file(os.path.join(pwd, 'npp_mask', shapefile_path))

    nodata = None
    mean = list()
    path = './mean_npp/'
    for file in sorted(os.listdir(path)):
        with rasterio.open(os.path.join(path, file)) as src:
            nodata = src.nodata
            # 应用掩膜
            out_image, out_transform = mask(src, mask_gdf.geometry, crop=True)
            out_image = out_image[0]  # 读取第一个波段
            image_except_nodata = out_image[out_image != nodata]
            mean.append(image_except_nodata.mean())

    # 统计 NoData 之外的值
    mean_npp = np.array(mean) / 10
    results[mask_chinese] = mean_npp

years = np.arange(2002, 2002 + len(mean_npp))

fig, ax1 = plt.subplots(figsize=(10, 8))

# 定义柔和的颜色和标记样式
colors = ['#6baed6', '#fd8d3c', '#74c476']  # 柔和的蓝色、橙色、绿色
markers = ['o', 's', '^']

# 用于记录已使用的文本位置，避免重叠
used_positions = []

# 绘制每个掩膜的年NPP值和趋势线
for i, (mask_chinese, mean_npp) in enumerate(results.items()):
    ax1.plot(years, mean_npp, marker=markers[i], label='{}年NPP值'.format(mask_chinese), color=colors[i])
    z = np.polyfit(years, mean_npp, 1)
    p = np.poly1d(z)
    ax1.plot(years, p(years), linestyle='--', label='{}趋势线'.format(mask_chinese), color=colors[i])
    
    # 动态调整文本位置，放置在曲线的最后一个点附近
    x_text = years[-1]  # 最后一个年份
    y_text = mean_npp[-1]  # 最后一个年份对应的值

    # 检查是否与已有文本位置重叠，若重叠则调整 y_text
    print(y_text)
    while any(abs(y_text - pos) < 40 for pos in used_positions):  # 0.5 为最小间距
        print('重叠！！')
        y_text -= 40  # 向上调整位置

    # 记录当前文本位置
    used_positions.append(y_text)

    ax1.text(x_text + 0.5, y_text, f'{mask_chinese}\ny={z[0]:.4f}x+{z[1]:.2f}\n$R^2$={np.corrcoef(years, mean_npp)[0, 1]**2:.8f}', 
             fontsize=10, color=colors[i], verticalalignment='center')

# 设置只保留左边和下方的轴
ax1.spines['top'].set_visible(False)  # 隐藏顶部边框
ax1.spines['right'].set_visible(False)  # 隐藏右侧边框

ax1.set_ylabel('NPP变化 (gc·m$^{-2}$·a$^{-1}$)')
ax1.set_xlabel('年份')
ax1.legend()
ax1.set_xticks(years)  # 设置年份刻度
ax1.set_xticklabels(years, rotation=45)  # 以整数形式显示年份，并旋转以防止重叠

plt.tight_layout()
plt.savefig('npp_trend_combined.png', dpi=300)
plt.show()